from pymilvus import (
    connections,
    Collection,
    CollectionSchema,
    FieldSchema,
    DataType
)
from sentence_transformers import SentenceTransformer
import uuid

class MilvusHelper:
    def __init__(self, host='localhost', port='19530', collection_name='article_paragraph'):
        self.host = host
        self.port = port
        self.collection_name = collection_name
        self.connect()
        self.model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
        self.collection = self.init_collection()

    def connect(self):
        connections.connect(alias='default', host=self.host, port=self.port)
        print(f"Connected to Milvus at {self.host}:{self.port}")

    def init_collection(self):
        paragraph_id = FieldSchema(name="paragraph_id", dtype=DataType.INT64, is_primary=True)
        article_id = FieldSchema(name="article_id", dtype=DataType.INT64)
        word_count = FieldSchema(name="word_count", dtype=DataType.INT64)
        paragraph_vector = FieldSchema(name="paragraph_vector", dtype=DataType.FLOAT_VECTOR, dim=768)
        
        schema = CollectionSchema(fields=[paragraph_id, article_id, word_count, paragraph_vector], description="Paragraph search")
        
        if not Collection(self.collection_name).exists:
            collection = Collection(name=self.collection_name, schema=schema, using='default', shards_num=8, consistency_level="Strong")
            print(f"Collection {self.collection_name} created.")
        else:
            collection = Collection(name=self.collection_name)
            print(f"Collection {self.collection_name} already exists.")
        
        return collection

    def insert_data(self, paragraph_id, article_id, sentence):
        embeddings = self.model.encode(sentence)
        data = [
            [paragraph_id],  # paragraph_id
            [article_id],    # article_id
            [0],             # dummy word_count
            [embeddings],    # paragraph_vector
        ]
        self.collection.insert(data)
        print("Data inserted.")

if __name__ == "__main__":
    milvus_helper = MilvusHelper()
    # Example usage
    paragraph_id = int(uuid.uuid4().int & (1<<64)-1)  # Example paragraph_id, using UUID for uniqueness
    article_id = 33333  # Example article_id
    sentence = "Hello, world!"
    milvus_helper.insert_data(paragraph_id, article_id, sentence)
